"""检索增强问答链：结合向量检索和LLM生成回答"""
from typing import Dict, List, Optional
import time
from src.vector_db.base_vector_db import BaseVectorDB  # 你的向量数据库基类
from src.llm.base_llm import BaseLLM  # 你的LLM基类
from src.llm.llm_factory import LLMFactory  # 复用你的LLM工厂
from config.app_config import get_app_config


class RetrievalQA:
    def __init__(
        self,
        vector_db: BaseVectorDB,  # 向量数据库实例（如ChromaDB）
        llm: Optional[BaseLLM] = None,  # LLM实例（可选，默认用配置中的模型）
        top_k: int = 3  # 默认检索前3个相关文档
    ):
        self.vector_db = vector_db
        self.top_k = top_k
        # 初始化LLM（默认用配置中的模型，通过工厂创建）
        self.llm = llm or LLMFactory.get_llm()
        self.config = get_app_config()

    def _build_prompt(self, question: str, context: str) -> str:
        """构造提示词：将问题和检索到的上下文组合"""
        return f"""
        基于以下上下文回答用户的问题，不要编造信息。如果上下文没有相关内容，直接说“没有找到相关信息”。
        
        上下文：
        {context}
        
        问题：{question}
        回答：
        """

    def run(self, question: str, streaming: bool = False) -> Dict[str, any]:
        """
        执行检索增强问答

        :param question: 用户的问题
        :param streaming: 是否流式返回回答（用于实时展示）
        :return: 包含回答、来源、处理时间的字典
        """
        start_time = time.time()

        # 1. 检索相关文档（调用向量数据库的相似性搜索）
        retrieved_docs = self.vector_db.similarity_search(
            query=question,
            k=self.top_k
        )
        # 提取文档内容和元数据（适配你的向量库返回格式）
        context = "\n\n".join([doc["page_content"] for doc in retrieved_docs])
        sources = [
            {
                "doc_id": doc["metadata"].get("doc_id", ""),
                "file_name": doc["metadata"].get("file_name", ""),
                "content": doc["page_content"][:100] + "..."  # 显示前100字符
            }
            for doc in retrieved_docs
        ]

        # 2. 调用LLM生成回答
        prompt = self._build_prompt(question, context)
        if streaming:
            # 流式返回（需LLM支持，如ollama的stream模式）
            answer = self.llm.stream(prompt)
        else:
            # 一次性返回
            answer = self.llm.generate(prompt)

        # 3. 计算处理时间并返回结果
        processing_time = round(time.time() - start_time, 2)
        return {
            "answer": answer,
            "sources": sources,
            "processing_time": processing_time
        }